//
// Created by cybaster on 6/5/25.
//

#pragma once

#include <mutex>

#include <ocs2_centroidal_model/CentroidalModelRbdConversions.h>
#include <ocs2_mpc/MPC_MRT_Interface.h>
#include <ocs2_mpc/SystemObservation.h>

#include <nav_msgs/Odometry.h>
#include <ros/ros.h>
#include <sensor_msgs/Imu.h>
#include <sensor_msgs/JointState.h>
#include <std_srvs/SetBool.h>

#include "ocs2_humanoid_robot/LeggedRobotInterface.h"

using namespace ocs2;
using namespace legged_robot;

class HumanoidRobotController {
public:
  explicit HumanoidRobotController(LeggedRobotInterface &leggedRobotInterface);

  virtual ~HumanoidRobotController() = default;

  virtual void launchNodes(ros::NodeHandle &nodeHandle);

protected:
  virtual vector_t calculateJointAccAndTorque(vector_t optimizedState,
                                        vector_t optimizedInput,
                                        size_t plannerMode);

  LeggedRobotInterface *interface_;
  CentroidalModelInfo info_;

  scalar_t timestep;
  vector_t rbdState_;

private:
  void setupMpc();

  void setupMrt();

  void start();

  void stop();

  void update();

  void updateState();

  void jointStateCallback(const sensor_msgs::JointState::ConstPtr &msg);

  void imuCallback(const sensor_msgs::Imu::ConstPtr &msg);

  void odometryCallback(const nav_msgs::Odometry::ConstPtr &msg);

  void timerCallback(const ros::TimerEvent &event);

  bool resetCallback(std_srvs::SetBool::Request &req,
                     std_srvs::SetBool::Response &res);

  std::string robotName_;

  std::shared_ptr<CentroidalModelRbdConversions> rbdConversions_;

  ros::Subscriber jointStateSubscriber_;
  ros::Subscriber imuSubscriber_;
  ros::Subscriber odometrySubscriber_;

  ros::Publisher mpcObservationPublisher_;
  ros::Publisher mpcCtrlPublisher_;

  ros::Timer timer_;
  ros::ServiceServer resetServer_;

  vector_t jointPos_, jointVel_;
  vector3_t pos_;
  Eigen::Quaternion<scalar_t> quat_;
  vector3_t angularVelLocal_, linearVelLocal_;

  std::mutex jointMutex_;
  std::mutex imuMutex_;
  std::mutex odometryMutex_;

  std::shared_ptr<MPC_BASE> mpc_;
  std::shared_ptr<MPC_MRT_Interface> mpcMrtInterface_;
  std::thread mpcThread_;
  std::atomic_bool controllerRunning_{}, mpcRunning_{};
  SystemObservation currentObservation_;
  size_t plannerMode_;
};
